import os
from PIL import Image, ImageDraw
import numpy as np
from scipy.ndimage import label, find_objects

def find_red_rectangles(effect_image, threshold=200):
    img_np = np.array(effect_image.convert("RGB"))
    red_mask = (img_np[:, :, 0] > threshold) & (img_np[:, :, 1] < 100) & (img_np[:, :, 2] < 100)
    labeled_array, num_features = label(red_mask)
    boxes = []

    for idx in range(1, num_features + 1):
        slices = find_objects(labeled_array == idx)[0]
        top, bottom = slices[0].start, slices[0].stop
        left, right = slices[1].start, slices[1].stop
        boxes.append((left, top, right, bottom))

    return boxes

def apply_transparent_mask_and_unmask(image, boxes, mask_color=(0, 0, 0, 128)):
    mask = Image.new("RGBA", image.size, mask_color)
    draw = ImageDraw.Draw(mask)
    draw.rectangle([0, 0, image.size[0], image.size[1]], fill=mask_color)

    for box in boxes:
        draw.rectangle(box, fill=(0, 0, 0, 0))

    result_image = Image.alpha_composite(image.convert("RGBA"), mask)
    return result_image

def process_folder(folder_path, output_folder):
    # 创建输出文件夹
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # 获取文件夹内的所有图片
    image_files = [f for f in os.listdir(folder_path) if f.endswith(('.jpg', '.jpeg', '.png'))]

    # 组织文件，按组别
    grouped_files = {}
    for file in image_files:
        parts = file.split("-")
        group = parts[0]
        role = parts[1].split(".")[0]  # 去掉文件扩展名
        if group not in grouped_files:
            grouped_files[group] = {}
        grouped_files[group][role] = file

    # 对每组找到原图和效果图并处理
    for group, files in grouped_files.items():
        if '1' in files and '2' in files:
            original_path = os.path.join(folder_path, files['1'])
            effect_path = os.path.join(folder_path, files['2'])

            original_image = Image.open(original_path)
            effect_image = Image.open(effect_path)

            # 找到效果图中的红色矩形
            red_boxes = find_red_rectangles(effect_image)

            # 应用透明蒙版并去除指定区域
            processed_image = apply_transparent_mask_and_unmask(original_image, red_boxes)

            # 保存处理后的图像
            output_file = f"{group}-processed.png"
            processed_image.save(os.path.join(output_folder, output_file))

            print(f"Processed {original_path} and {effect_path} into {output_folder}/{output_file}")

# 示例使用
if __name__ == "__main__":
    folder_path = r"C:\Users\Administrator\Desktop\test"  # 替换为你的文件夹路径
    output_folder = "./test"  # 替换为输出文件夹路径
    process_folder(folder_path, output_folder)
